{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sb_auto_header",
    "tags": [
     "sb_auto_header"
    ]
   },
   "source": [
    "<!-- This cell is automatically updated by tools/tutorial-cell-updater.py -->\n",
    "<!-- The contents are initialized from tutorials/notebook-header.md -->\n",
    "\n",
    "[<img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>](https://colab.research.google.com/github/speechbrain/speechbrain/blob/develop/docs/tutorials/advanced/hyperparameter-optimization.ipynb)\n",
    "to execute or view/download this notebook on\n",
    "[GitHub](https://github.com/speechbrain/speechbrain/tree/develop/docs/tutorials/advanced/hyperparameter-optimization.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "GxuFVVbnJb8d"
   },
   "source": [
    "# Hyperparameter Optimization\n",
    "\n",
    "Many of the speech processing tasks implemented as part of the SpeechBrain project rely on the careful selection of hyperparameters, such as:\n",
    "\n",
    "* The number of layers\n",
    "* Normalization\n",
    "* Hidden layer dimensions\n",
    "* Weights within cost functions\n",
    "* etc\n",
    "\n",
    "Selecting such hyperparameters by hand can be tedious. This tutorial will show how to use the automated hyperparameter optimization techniques implemented as part of the [Oríon](https://github.com/Epistimio/orion) project to automatically fit hyperparameters in a systematic way."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JnRSS3jGxVsk"
   },
   "source": [
    "## Prerequisites\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C0crB6tRbRFN"
   },
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m2X6fF7dbCRz"
   },
   "outputs": [],
   "source": [
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vxr8nZ7bbP15"
   },
   "source": [
    "### Install SpeechBrain\n",
    "\n",
    "SpeechBrain can be downloaded from the GitHub repository listed below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "KSOTkqPsJXxZ"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "# Installing SpeechBrain via pip\n",
    "BRANCH = 'develop'\n",
    "!python -m pip install git+https://github.com/speechbrain/speechbrain.git@$BRANCH\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BTlxRL6aNtCp"
   },
   "source": [
    "### Dependency Fixes\n",
    "\n",
    "PyYAML 6.0 is not backwards-compatible, a 5.x version is needed to support HyperPyYAML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EX8eL65Q6B4D"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install pyyaml==5.4.1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iakT5Lq3xvSU"
   },
   "source": [
    "### Install Oríon\n",
    "Oríon can be installed using `pip` or `conda`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "09WKjQpoyHxl"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "!pip install orion[profet]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CowQ-8Ia7pWo"
   },
   "outputs": [],
   "source": [
    "from speechbrain.utils import hpopt as hp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hEcF9W_cyxpv"
   },
   "source": [
    "## Update the Recipe to Support Hyperparameter Optimization\n",
    "\n",
    "SpeechBrain comes with a convenience wrapper called `hpopt`, which is capable of reporting objective values to Orion or to other tools.\n",
    "\n",
    "For a complete example on how to implement it,\n",
    "\n",
    "1. Add the following import statement to the top of your recipe:\n",
    "\n",
    "```python\n",
    "from speechbrain.utils import hpopt as hp\n",
    "```\n",
    "\n",
    "2. Wrap the main code of your recipe in a hyperparameter optimization context. Set `objective_key` to the metric that Orion will optimize.\n",
    "\n",
    "    **Before**:\n",
    "\n",
    "  ```python\n",
    "  hparams_file, run_opts, overrides = sb.parse_arguments(sys.argv[1:])\n",
    "  \n",
    "  with open(hparams_file) as fin:\n",
    "      hparams = load_hyperpyyaml(fin, overrides)\n",
    "  \n",
    "  ## ...\n",
    "  \n",
    "  spk_id_brain = SpkIdBrain(\n",
    "      modules=hparams[\"modules\"],\n",
    "      opt_class=hparams[\"opt_class\"],\n",
    "      hparams=hparams,\n",
    "      run_opts=run_opts,\n",
    "      checkpointer=hparams[\"checkpointer\"],\n",
    "  )\n",
    "  \n",
    "  # The `fit()` method iterates the training loop, calling the methods\n",
    "  # necessary to update the parameters of the model. Since all objects\n",
    "  # with changing state are managed by the Checkpointer, training can be\n",
    "  # stopped at any point, and will be resumed on next call.\n",
    "  spk_id_brain.fit(\n",
    "      epoch_counter=spk_id_brain.hparams.epoch_counter,\n",
    "      train_set=datasets[\"train\"],\n",
    "      valid_set=datasets[\"valid\"],\n",
    "      train_loader_kwargs=hparams[\"dataloader_options\"],\n",
    "      valid_loader_kwargs=hparams[\"dataloader_options\"],\n",
    "  )\n",
    "\n",
    "  ```\n",
    "\n",
    "  **After**:\n",
    "\n",
    "    ```python\n",
    "    with hp.hyperparameter_optimization(objective_key=\"error\") as hp_ctx: # <-- Initialize the context\n",
    "        hparams_file, run_opts, overrides = hp_ctx.parse_arguments(sys.argv[1:]) # <-- Replace sb with hp_ctx\n",
    "\n",
    "        with open(hparams_file) as fin:\n",
    "            hparams = load_hyperpyyaml(fin, overrides)\n",
    "\n",
    "        ## ...\n",
    "\n",
    "            spk_id_brain = SpkIdBrain(\n",
    "                modules=hparams[\"modules\"],\n",
    "                opt_class=hparams[\"opt_class\"],\n",
    "                hparams=hparams,\n",
    "                run_opts=run_opts,\n",
    "                checkpointer=hparams[\"checkpointer\"],\n",
    "            )\n",
    "\n",
    "            # The `fit()` method iterates the training loop, calling the methods\n",
    "            # necessary to update the parameters of the model. Since all objects\n",
    "            # with changing state are managed by the Checkpointer, training can be\n",
    "            # stopped at any point, and will be resumed on next call.\n",
    "            spk_id_brain.fit(\n",
    "                epoch_counter=spk_id_brain.hparams.epoch_counter,\n",
    "                train_set=datasets[\"train\"],\n",
    "                valid_set=datasets[\"valid\"],\n",
    "                train_loader_kwargs=hparams[\"dataloader_options\"],\n",
    "                valid_loader_kwargs=hparams[\"dataloader_options\"],\n",
    "            )\n",
    "    ```\n",
    "\n",
    "3. Add code to report the stats\n",
    "\n",
    "  e.g. in `on_stage_end` when `stage == sb.Stage.VALID`\n",
    "\n",
    "  ```python\n",
    "hp.report_result(stage_stats)\n",
    "```\n",
    "\n",
    "  The **last** result reported through this function will be reported for hyperparameter optimization.\n",
    "\n",
    "  The key specified in **objective_key** parameter needs to be present in the dictionary passed to `report_result`.\n",
    "\n",
    "4. Add the following lines in your main hyperparameter file `train.yaml`:\n",
    "```yaml\n",
    "hpopt_mode: null\n",
    "hpopt: null\n",
    "```\n",
    "\n",
    "5. **Optional**: Create a separate YAML file overriding any hyperparameters to be used during hyperparameter optimization that are **different** from the ones used during regular training **other than** the ones being fitted. A typical approach would reduce the number of epochs and the number of training samples.\n",
    "\n",
    "  This step can be omitted if the number of parameters being overridden is small. In this case, they can be passed on the command line instead.\n",
    "\n",
    "  Example:\n",
    "\n",
    "  `hpopt.yaml`:\n",
    "  ```yaml\n",
    "  number_of_epochs: 1\n",
    "  ckpt_enable: false\n",
    "  ```\n",
    "6. ❗ **Important**: Most recipes use a checkpointer to save snapshots of the model after each epoch (or on a custom schedule) to ensure that training can be resumed if it is interrupted. During hyperparameter optimization, this can cause issues because if the model's architecture (e.g. the number of layers, neurons per layer, etc) changes from one set of hyperparamter values to the next, an attempt to restore a checkpoint will fail.\n",
    "\n",
    "  One possible solution is to make the run of the checkpointer conditional and to disable it in `hpopt.yaml`\n",
    "\n",
    "  __Before__:\n",
    "  ```python\n",
    "  self.checkpointer.save_and_keep_only(meta=stats, min_keys=[\"error\"])\n",
    "  ```\n",
    "  __After__:\n",
    "  ```python\n",
    "  if self.hparams.ckpt_enable:\n",
    "      self.checkpointer.save_and_keep_only(meta=stats, min_keys=[\"error\"])\n",
    "  ```\n",
    "\n",
    "  An alternative strategy is to reconfigure the checkpointer to save each run in a separate directory. For this scenario, the hyperparameter optimization wrapper can supply a variable named trial_id, which can be interpolated into the output path.\n",
    "\n",
    "  Given below is an example of this strategy:\n",
    "\n",
    "  `hpopt.yaml`:\n",
    "\n",
    "  ```yaml\n",
    "  number_of_epochs: 1\n",
    "  ckpt_enable: False\n",
    "  trial_id: hpopt\n",
    "  output_folder: !ref ./results/speaker_id/<trial_id>\n",
    "  ```\n",
    "\n",
    "  `train.yaml`:\n",
    "\n",
    "  ```yaml\n",
    "  # ...\n",
    "  save_folder: !ref <output_folder>/save\n",
    "  # ...\n",
    "  checkpointer: !new:speechbrain.utils.checkpoints.Checkpointer\n",
    "      checkpoints_dir: !ref <save_folder> #<-- will contain trial_id\n",
    "      recoverables:\n",
    "          embedding_model: !ref <embedding_model>\n",
    "          classifier: !ref <classifier>\n",
    "          normalizer: !ref <mean_var_norm>\n",
    "          counter: !ref <epoch_counter>\n",
    "  ```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TwbiUJLWSICU"
   },
   "source": [
    "## Perform the Hyperparameter Search\n",
    "\n",
    "### Choose and Prepare Hyperparameters\n",
    "\n",
    "Choose the hyperparameters that you would like to optimize using Orion out of the\n",
    "ones available in your hyperparameter file. The hyperparameters need to be\n",
    "available at the top level in order for it to be fitted using this technique.\n",
    "\n",
    "Consider the following sample file:\n",
    "\n",
    "```yaml\n",
    "dropout: 0.1\n",
    "n_mels: 80\n",
    "encoder: !new:speechbrain.lobes.models.mymodel.MyModel\n",
    "    input_shape: [null, null, !ref <n_mels>]\n",
    "    dropout: !ref <dropout>\n",
    "    cnn_blocks: 3\n",
    "```\n",
    "\n",
    "In the above file, `n_mels` and `dropout` are available for optimization, but `cnn_blocks` is not.\n",
    "\n",
    "To make `cnn_blocks` available for optimization, modify it as follows:\n",
    "\n",
    "```yaml\n",
    "dropout: 0.1\n",
    "n_mels: 80\n",
    "cnn_blocks: 3 # <-- Define at the top level\n",
    "encoder: !new:speechbrain.lobes.models.mymodel.MyModel\n",
    "    input_shape: [null, null, !ref <n_mels>]\n",
    "    dropout: !ref <dropout>\n",
    "    cnn_blocks: !ref <cnn_blocks> # <-- Introduce a reference\n",
    "```\n",
    "\n",
    "### Configure Orion\n",
    "Create a `.yaml` file with the configuration for the Orion algorithm to be used.\n",
    "\n",
    "Given below is an example:\n",
    "```yaml\n",
    "experiment:\n",
    "    max_trials: 1000\n",
    "    max_broken: 1000\n",
    "    algorithms:\n",
    "        tpe:\n",
    "            seed: 42\n",
    "            n_initial_points: 5\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E_kGgrp_Z35H"
   },
   "outputs": [],
   "source": [
    "config_file_content = \"\"\"\n",
    "experiment:\n",
    "    max_trials: 3\n",
    "    max_broken: 1\n",
    "    algorithms:\n",
    "        tpe:\n",
    "            seed: 42\n",
    "            n_initial_points: 5\n",
    "\"\"\"\n",
    "config_path = os.path.expanduser(\"~/config\")\n",
    "if not os.path.exists(config_path):\n",
    "    os.mkdir(config_path)\n",
    "\n",
    "config_file_path = os.path.join(config_path, \"orion-speaker-id.yaml\")\n",
    "with open(config_file_path, \"w\") as config_file:\n",
    "    print(config_file_content, file=config_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5Kt-blyWadDV"
   },
   "source": [
    "For more information on the available algorithms, please take a look at  the [Orion Repository](https://github.com/Epistimio/orion/tree/develop/src/orion/algo)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xVyU61qXaCox"
   },
   "source": [
    "### Define the Search Space\n",
    "Write a shell script calling Orion defining a search space\n",
    "\n",
    "Example:\n",
    "```sh\n",
    "#!/bin/bash\n",
    "HPOPT_EXPERIMENT_NAME=speaker-id\n",
    "HPOPT_CONFIG_FILE=$HOME/config/orion-speaker-id.yaml\n",
    "orion hunt -n $HPOPT_EXPERIMENT_NAME -c $HPOPT_CONFIG_FILE python train.py hparams/$HPARAMS \\\n",
    "    --hpopt hpopt.yaml \\\n",
    "    --hpopt_mode orion \\\n",
    "    --emb_dim~\"choices([128,256,512,768,1024])\" \\\n",
    "    --tdnn_channels~\"choices([128,256,512,768,1024])\"\n",
    "```\n",
    "\n",
    "Replace `--hpopt hpopt.yaml` with `--hpopt=True` if you are not using the additional  `hpopt.yaml` file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xCw_DZYOZ7r0"
   },
   "source": [
    "Consider running the standalone example below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lJup9mNnYw_0"
   },
   "outputs": [],
   "source": [
    "%env PYTHONPATH=/env/python:/content/speechbrain/\n",
    "%cd /content/speechbrain/templates/hyperparameter_optimization_speaker_id\n",
    "!orion hunt -n speaker-id -c $HOME/config/orion-speaker-id.yaml python train.py train.yaml \\\n",
    "  --hpopt hpopt.yaml \\\n",
    "  --hpopt_mode orion \\\n",
    "  --emb_dim~\"choices([128,256,512,768,1024])\" \\\n",
    "  --tdnn_channels~\"choices([128,256,512,768,1024])\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FntqoHyK8Llt"
   },
   "source": [
    "## Inspecting Results\n",
    "\n",
    "Use the `orion info` command to inspect the results of hyperparameter fitting.\n",
    "\n",
    "The tool will output basic statistics about the hyperparameter fitting experiment, including the number of runs completed, the objective value for the best trial and the hyperparameter values corresponding to that run.\n",
    "\n",
    "In the example below, the best objective achieved value is shown under **evaluation**, and the corresponding hyperparameter values are shown under **params**.\n",
    "\n",
    "```\n",
    "Stats\n",
    "=====\n",
    "completed: False\n",
    "trials completed: 4\n",
    "best trial:\n",
    "  id: c1a71e0988d70005302ab655d7e391d3\n",
    "  evaluation: 0.2384105920791626\n",
    "  params:\n",
    "    /emb_dim: 128\n",
    "    /tdnn_channels: 128\n",
    "start time: 2021-11-14 21:01:12.760704\n",
    "finish time: 2021-11-14 21:13:25.043336\n",
    "duration: 0:12:12.282632\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lVVUBZgYBNQP"
   },
   "outputs": [],
   "source": [
    "!orion info --name speaker-id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ys-hUnmgSsFf"
   },
   "source": [
    "## Hyperparameter Optimization at Scale\n",
    "\n",
    "### Multiple GPUs\n",
    "Since Orion simply wraps the execution of the training script and launches it for each set of hyperparameters using the OS shell, training scripts that support Data-Parallel (DP) or Distributed Data Parallel (DDP) execution can be used with hyperparameter fitting without modification.\n",
    "\n",
    "For information on how to set up DP/DDP experiments, refer to the [SpeechBrain documentation](https://speechbrain.readthedocs.io/en/latest/multigpu.html#) and the [Multi-GPU Considerations](https://speechbrain.readthedocs.io/en/latest/multigpu.html) tutorial.\n",
    "\n",
    "### Parallel or Distributed Oríon\n",
    "\n",
    "Oríon itself provide support for parallel and distributed hyperparameter fitting.\n",
    "\n",
    "To use multiple parallel workers on a single node, pass the `--n-workers` parameter to the Oríon CLI.\n",
    "\n",
    "The example below will start the experiment with three workers:\n",
    "```shell\n",
    "orion hunt -n $HPOPT_EXPERIMENT_NAME -c $HOPT_CONFIG_FILE --n-workers 3 python train.py hparams/$HPARAMS \\\n",
    "    --hpopt hpopt.yaml \\\n",
    "    --hpopt_mode orion \\\n",
    "    --emb_dim~\"choices([128,256,512,768,1024])\" \\\n",
    "    --tdnn_channels~\"choices([128,256,512,768,1024])\"\n",
    "\n",
    "```\n",
    "\n",
    "For more advanced scenarios, including distributed hyperparameter fittig on multiple nodes, refer to the [Parallel Workers](https://orion.readthedocs.io/en/stable/user/parallel.html]) page in Oríon's official documentation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sb_auto_footer",
    "tags": [
     "sb_auto_footer"
    ]
   },
   "source": [
    "## Citing SpeechBrain\n",
    "\n",
    "If you use SpeechBrain in your research or business, please cite it using the following BibTeX entry:\n",
    "\n",
    "```bibtex\n",
    "@misc{speechbrainV1,\n",
    "  title={Open-Source Conversational AI with {SpeechBrain} 1.0},\n",
    "  author={Mirco Ravanelli and Titouan Parcollet and Adel Moumen and Sylvain de Langen and Cem Subakan and Peter Plantinga and Yingzhi Wang and Pooneh Mousavi and Luca Della Libera and Artem Ploujnikov and Francesco Paissan and Davide Borra and Salah Zaiem and Zeyu Zhao and Shucong Zhang and Georgios Karakasidis and Sung-Lin Yeh and Pierre Champion and Aku Rouhe and Rudolf Braun and Florian Mai and Juan Zuluaga-Gomez and Seyed Mahed Mousavi and Andreas Nautsch and Xuechen Liu and Sangeet Sagar and Jarod Duret and Salima Mdhaffar and Gaelle Laperriere and Mickael Rouvier and Renato De Mori and Yannick Esteve},\n",
    "  year={2024},\n",
    "  eprint={2407.00463},\n",
    "  archivePrefix={arXiv},\n",
    "  primaryClass={cs.LG},\n",
    "  url={https://arxiv.org/abs/2407.00463},\n",
    "}\n",
    "@misc{speechbrain,\n",
    "  title={{SpeechBrain}: A General-Purpose Speech Toolkit},\n",
    "  author={Mirco Ravanelli and Titouan Parcollet and Peter Plantinga and Aku Rouhe and Samuele Cornell and Loren Lugosch and Cem Subakan and Nauman Dawalatabad and Abdelwahab Heba and Jianyuan Zhong and Ju-Chieh Chou and Sung-Lin Yeh and Szu-Wei Fu and Chien-Feng Liao and Elena Rastorgueva and François Grondin and William Aris and Hwidong Na and Yan Gao and Renato De Mori and Yoshua Bengio},\n",
    "  year={2021},\n",
    "  eprint={2106.04624},\n",
    "  archivePrefix={arXiv},\n",
    "  primaryClass={eess.AS},\n",
    "  note={arXiv:2106.04624}\n",
    "}\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [
    {
     "file_id": "14Lh3BPve730S8NhbxypYeTrq16R32Xp5",
     "timestamp": 1636476890780
    }
   ]
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
